import torchvision
import sys
import numpy as np
import torch
import matplotlib.pyplot as plt
import cv2
from segment_anything import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor

def show_anns(anns):
    if len(anns) == 0:
        return
    sorted_anns = sorted(anns, key=(lambda x: x['area']), reverse=True)
    ax = plt.gca()
    ax.set_autoscale_on(False)
    polygons = []
    color = []
    for ann in sorted_anns:
        m = ann['segmentation']
        img = np.ones((m.shape[0], m.shape[1], 3))
        color_mask = np.random.random((1, 3)).tolist()[0]
        for i in range(3):
            img[:,:,i] = color_mask[i]
        ax.imshow(np.dstack((img, m*0.35)))


image = cv2.imread('uploads/a.jpg')
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

sam_checkpoint = "/apps/segment-anything/models/sam_vit_h_4b8939.pth"

device = "cuda"
model_type = "default"

sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
sam.to(device=device)

input_point = np.array([[523, 159]])
input_label = np.array([1])

predictor = SamPredictor(sam)
predictor.set_image(image) # embedding操作
masks, scores, logits = predictor.predict(
   point_coords=input_point,
   point_labels=input_label,
   multimask_output=True,)

# mask_generator = SamAutomaticMaskGenerator(
#   sam,
#   points_per_side=64, #默认32
#   pred_iou_thresh=0.8, #默认0.98
#   stability_score_thresh=0.9, #默认0.95
#   crop_n_layers=1,
#   crop_n_points_downscale_factor=2,
#   min_mask_region_area=10,  # Requires open-cv to run post-processing
# )
# masks = mask_generator.generate(image)

plt.figure(figsize=(20,20))
plt.imshow(image)
# show_anns(masks)
plt.axis('off')
plt.savefig('uploads/a_out.jpg')
# plt.show()